from celery import Celery
import os
from model.aoi import Lot, Wafer, Image, Failure
from model.pmi import PMILot, PMIWafer, PMIImage, PMIDevice
import json
from weice_package import weice
from utils.wafer import get_wafer_map
from utils.file import rename_pmi
from datetime import datetime, timedelta
from utils.database import SessionLocal
from utils.msg_redis import MessageRedis
from decimal import Decimal
import time

# 创建 Celery 实例
celery_app = Celery('tasks', broker='amqp://guest@localhost//')

redis = MessageRedis()

@celery_app.task(argsrepr=True)
def detect_aoi_lot_celery(aoi_yolox_model, aoi_seg_model, cls_model,lot_folder="/data/lot/all_lot", distance=6, set_num_zhenhen=3, set_diagonal_length=30, set_area=25, set_per_area=0.1):
    print("detect_aoi_lot")
    print(lot_folder)

    # 后台任务只能启动一次
    redis.set_string("task", "start")

    with SessionLocal() as session:
    
        try:
            for folder in os.listdir(lot_folder):
                
                # 获取 lot 文件夹
                lotname = folder
                lot_dir = os.path.join(lot_folder, folder)

                lot = session.query(Lot).filter_by(lot_name=lotname).first()

                if lot:
                    print("1111111111111已存在数据库中: ", lotname)
                    continue
                else:
                    lot = Lot(lot_name=lotname, dir_path=lot_dir)
                    print("222222222222222222开始处理lot: ", lotname)
                    session.add(lot)
                    session.commit()

                    # 开始推理每个 wafer folder
                    for file in os.listdir(lot_dir):
                        # 获取 wafer 矩阵文件
                        if file.endswith('.txt'):
                            wafer_index = file.split("-")[1]
                            current_time = (datetime.utcnow() + timedelta(hours=8)).strftime('%Y-%m-%d %H:%M:%S')
                            json_data = {"msg": 'Lot %s Wafer %s 开始处理' % (lotname, wafer_index), "time": current_time, "status": "start", "level": "info"}
                            # redis.set_message("aoi", json_data)

                            wafer_folder = os.path.join(lot_dir, wafer_index)

                            wafer_txt = os.path.join(lot_dir, file)
                            print("aaaaaaaaaaaaaaaaaa开始处理wafer: ", wafer_index)

                            t1 = time.time()
                            predict_data = weice.run_aoi(aoi_yolox_model, aoi_seg_model, cls_model, wafer_folder, distance=distance, set_num_zhenhen=set_num_zhenhen, set_diagonal_length=set_diagonal_length, set_area=set_area, set_per_area=set_per_area)

                            t2 = time.time()
                            print("bbbbbbbbbbbbbbbbbbbb结束处理wafer: ", wafer_index)

                            print("推理耗时：", t2-t1)

                            # 生成 wafer map 矩形图
                            wafer_data = get_wafer_map(wafer_txt, wafer_folder, predict_data)

                            wafer = Wafer(wafer_id=wafer_index, lot_code=lot.lot_code, wafer_data=json.dumps(wafer_data))
                            session.add(wafer)
                            session.commit()

                            # 每张图片的推理结果
                            for image_predict in predict_data:
                                failure_id = 1
                                detection_info = {}
                                segmentation_info = {}
                                image_name = image_predict["filename"]

                                wuran_box = image_predict.get("wuran_box", [])
                                wuran_crop = image_predict.get("wuran_crop", [])

                                can_wuran_box = image_predict.get("can_wuran_box", [])

                                zhenhen_box = image_predict.get("zhenhen_box", [])
                                zhenhen_crop = image_predict.get("zhenhen_crop", [])
                                # zhenhen_num_error = image_predict.get("zhenhen_num_error", [])

                                madian_box = image_predict.get("madian_box", [])

                                # if zhenhen_num_error:
                                #     detection_info[4] = zhenhen_num_error
                                #     failure_id = 4

                                if zhenhen_box:
                                    detection_info[3] = zhenhen_box
                                    failure_id = 3
                                
                                if zhenhen_crop:
                                    segmentation_info[3] = zhenhen_crop
                                    failure_id = 3

                                if can_wuran_box:
                                    detection_info[5] = can_wuran_box
                                    failure_id = 5

                                if wuran_box:
                                    detection_info[2] = wuran_box
                                    failure_id = 2

                                if wuran_crop:
                                    segmentation_info[2] =  wuran_crop
                                    failure_id = 2

                                if madian_box:
                                    detection_info[6] =  madian_box
                                    failure_id = 6

                                file_path = os.path.join(lot_dir, wafer_index, image_name)

                                if detection_info:
                                    detection_info = json.dumps(detection_info)
                                else:
                                    detection_info = ""

                                if segmentation_info:
                                    segmentation_info=json.dumps(segmentation_info)
                                else:
                                    segmentation_info = ""

                                image = Image(
                                    wafer_code = wafer.wafer_code, 
                                    file_path = file_path,
                                    file_name = image_name,
                                    detection_info=detection_info,
                                    segmentation_info=segmentation_info,
                                    failure_id = failure_id
                                )
                                
                                session.add(image)

                            # wafer 处理完
                            current_time = (datetime.utcnow() + timedelta(hours=8)).strftime('%Y-%m-%d %H:%M:%S')
                            json_data = {"msg": 'Lot %s Wafer %s 处理完成' % (lotname, wafer_index), "time": current_time, "status": "end", "level": "info"}  
                            # redis.set_message("aoi", json_data)

                # lot 处理完
                current_time = (datetime.utcnow() + timedelta(hours=8)).strftime('%Y-%m-%d %H:%M:%S')
                json_data = {"msg": 'Lot %s 处理完成' % lotname, "time": current_time, "status": "completed", "level": "info"}  

                lot.status = "COMPLETED"
                session.commit()

                # redis.set_message("aoi", json_data)
                
            # current_time = (datetime.utcnow() + timedelta(hours=8)).strftime('%Y-%m-%d %H:%M:%S')
            # json_data = {"msg": '%s 下的所有 lot 处理完' % lot_folder, "time": current_time, "status": "completed"}             
            # socketio.emit('job update', json_data, namespace="/msg")

        except Exception as e:
            print("Exception: ", e)
            current_time = (datetime.utcnow() + timedelta(hours=8)).strftime('%Y-%m-%d %H:%M:%S')
            json_data = {"msg": 'Lot %s 处理出错' % lotname, "time": current_time, "level": "error"}
            # redis.set_message("aoi", json_data)
            # session.rollback()

        # 后台任务执行完成
        # redis.set_string("task", "end")